use std::path::Path;

use melior::{
    Context,
    ir::{
        Attribute, Block, BlockLike, BlockRef, Identifier, Location, Module, Region, RegionLike,
        Type, TypeLike, Value, ValueLike,
        attribute::{
            DenseElementsAttribute, FlatSymbolRefAttribute, FloatAttribute, StringAttribute,
            TypeAttribute,
        },
        operation::{OperationBuilder, OperationLike},
        r#type::{FunctionType, RankedTensorType},
    },
};
use thiserror::Error;

use crate::{
    parser::Parser,
    rust_side::{
        ScopedHashTable, ScopedHashTableScope, mlir_declare, mlir_lookup, mlir_set_private,
    },
};

#[derive(Debug, Error)]
pub enum CompileError {
    #[error("{0}")]
    ParseError(#[from] crate::parser::ParseError),
    #[error("{0}")]
    MeliorError(#[from] melior::Error),
    #[error("definition of {0} not found")]
    NotFound(String),
    #[error("{0}")]
    Utf8Error(#[from] std::str::Utf8Error),
    #[error("{0}")]
    IntoStringError(#[from] std::ffi::IntoStringError),
}

pub type Result<T> = std::result::Result<T, CompileError>;

pub struct Compiler {
    scoped_hash_table: ScopedHashTable,
}

impl Compiler {
    pub fn new() -> Self {
        Self {
            scoped_hash_table: ScopedHashTable::new(),
        }
    }

    pub fn compile<'c, P: AsRef<Path>>(
        &mut self,
        context: &'c Context,
        path: P,
    ) -> Result<Module<'c>> {
        let mut parser = Parser::new(path)?;
        let m = parser.parse()?;
        self.compile_module(context, &m)
    }

    fn compile_module<'c>(
        &self,
        context: &'c Context,
        m: &crate::ast::Module,
    ) -> Result<Module<'c>> {
        let location = Location::unknown(context);

        let module = Module::new(location);
        for func in &m.functions {
            self.compile_function(context, &module, func)?;
        }
        assert!(module.as_operation().verify());
        Ok(module)
    }

    fn compile_function(
        &self,
        context: &Context,
        module: &Module,
        func: &crate::ast::Function,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let _scope = ScopedHashTableScope::new(&self.scoped_hash_table);

        let mut func_type: FunctionType = self.compile_prototype(context, &func.proto)?;

        let mut arguments = Vec::new();
        for index in 0..func_type.input_count() {
            arguments.push((func_type.input(index)?, location));
        }
        let block = Block::new(&arguments);
        let body = Region::new();
        let block_ref = body.append_block(block);
        for index in 0..block_ref.argument_count() {
            let arg = block_ref.argument(index).unwrap();
            let name = &func.proto.args[index].name;
            mlir_declare(&self.scoped_hash_table, name, arg.into());
        }

        self.compile_expr_list(context, block_ref, &func.body)?;

        let return_op_ref = match block_ref.terminator() {
            Some(terminator) => {
                (terminator.name().as_string_ref().as_str()? == "toy.return").then_some(terminator)
            }
            None => None,
        };

        let sym_name = StringAttribute::new(context, &func.proto.name);
        match return_op_ref {
            Some(return_op_ref) => {
                let mut inputs = Vec::new();
                for index in 0..func_type.input_count() {
                    inputs.push(func_type.input(index).unwrap());
                }

                let results = return_op_ref
                    .operands()
                    .map(|operand| operand.r#type())
                    .collect::<Vec<Type>>();

                func_type = FunctionType::new(context, &inputs, &results);
            }
            None => {
                block_ref.append_operation(OperationBuilder::new("toy.return", location).build()?);
            }
        }

        let function_type = TypeAttribute::new(func_type.into());
        let func_op = OperationBuilder::new("toy.func", location)
            .add_attributes(&[
                (Identifier::new(context, "sym_name"), sym_name.into()),
                (
                    Identifier::new(context, "function_type"),
                    function_type.into(),
                ),
            ])
            .add_regions([body])
            .build()?;
        let func_op_ref = module.body().append_operation(func_op.into());
        if &func.proto.name != "main" {
            mlir_set_private(func_op_ref);
        }

        Ok(())
    }

    fn compile_prototype<'c>(
        &self,
        context: &'c Context,
        prototype: &crate::ast::Prototype,
    ) -> Result<FunctionType<'c>> {
        let mut inputs = Vec::new();
        for _ in &prototype.args {
            inputs.push(self.get_type(context, &[])?);
        }

        let results = &[];
        Ok(FunctionType::new(context, &inputs, results))
    }

    fn get_type<'c>(&self, context: &'c Context, dimensions: &[u64]) -> Result<Type<'c>> {
        let f64_type = Type::float64(context);
        Ok(if dimensions.is_empty() {
            unsafe { Type::from_raw(mlir_sys::mlirUnrankedTensorTypeGet(f64_type.to_raw())) }
        } else {
            RankedTensorType::new(dimensions, f64_type, None).into()
        })
    }

    fn compile_expr_list(
        &self,
        context: &Context,
        block_ref: BlockRef,
        expr_list: &crate::ast::ExprList,
    ) -> Result<()> {
        let _scope = ScopedHashTableScope::new(&self.scoped_hash_table);

        for expr in &expr_list.exprs {
            self.compile_expr(context, block_ref, expr)?;
        }
        Ok(())
    }

    fn compile_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        expr: &crate::ast::Expr,
    ) -> Result<Option<Value<'c, 'c>>> {
        Ok(match expr {
            crate::ast::Expr::NumberExpr(number_expr) => {
                let value = self.compile_number_expr(context, block_ref, number_expr)?;
                Some(value)
            }
            crate::ast::Expr::LiteralExpr(literal_expr) => {
                let value = self.compile_literal_expr(context, block_ref, literal_expr)?;
                Some(value)
            }
            crate::ast::Expr::VariableExpr(variable_expr) => {
                let value = self.compile_variable_expr(context, block_ref, variable_expr)?;
                Some(value)
            }
            crate::ast::Expr::VarDeclExpr(var_decl_expr) => {
                self.compile_var_decl_expr(context, block_ref, var_decl_expr)?;
                None
            }
            crate::ast::Expr::ReturnExpr(return_expr) => {
                self.compile_return_expr(context, block_ref, return_expr)?;
                None
            }
            crate::ast::Expr::BinaryExpr(binary_expr) => {
                let value = self.compile_binary_expr(context, block_ref, binary_expr)?;
                Some(value)
            }
            crate::ast::Expr::CallExpr(call_expr) => {
                let value = self.compile_call_expr(context, block_ref, call_expr)?;
                Some(value)
            }
            crate::ast::Expr::PrintExpr(print_expr) => {
                self.compile_print_expr(context, block_ref, print_expr)?;
                None
            }
        })
    }

    fn compile_number_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        number_expr: &crate::ast::NumberExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let f64_type = Type::float64(context);
        let r#type = RankedTensorType::new(&[], f64_type, None);
        let values = vec![FloatAttribute::new(context, f64_type, number_expr.val).into()];
        let value = DenseElementsAttribute::new(r#type.into(), &values)?;
        let constant_op = OperationBuilder::new("toy.constant", location)
            .add_results(&[r#type.into()])
            .add_attributes(&[(Identifier::new(context, "value"), value.into())])
            .build()?;
        let constant_op_ref = block_ref.append_operation(constant_op);
        Ok(constant_op_ref.result(0).unwrap().into())
    }

    fn compile_literal_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        literal_expr: &crate::ast::LiteralExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let mut values = Vec::new();
        self.collect_literal_expr_values(literal_expr, &mut values)?;

        let results = &[self.get_type(context, &literal_expr.dims)?];
        let f64_type = Type::float64(context);
        let r#type = RankedTensorType::new(&literal_expr.dims, f64_type, None).into();
        let values = values
            .iter()
            .map(|value| FloatAttribute::new(context, f64_type, *value).into())
            .collect::<Vec<Attribute>>();
        let value = DenseElementsAttribute::new(r#type, &values)?;
        let constant_op = OperationBuilder::new("toy.constant", location)
            .add_results(results)
            .add_attributes(&[(Identifier::new(context, "value"), value.into())])
            .build()?;
        let constant_op_ref = block_ref.append_operation(constant_op);
        let result = constant_op_ref.result(0).unwrap().into();
        Ok(result)
    }

    fn compile_variable_expr<'c>(
        &self,
        _context: &'c Context,
        _block_ref: BlockRef<'c, 'c>,
        variable_expr: &crate::ast::VariableExpr,
    ) -> Result<Value<'c, 'c>> {
        match mlir_lookup(&self.scoped_hash_table, &variable_expr.name) {
            Some(value) => Ok(value),
            None => Err(CompileError::NotFound(variable_expr.name.clone())),
        }
    }

    fn compile_var_decl_expr(
        &self,
        context: &Context,
        block_ref: BlockRef,
        var_decl_expr: &crate::ast::VarDeclExpr,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let mut value = self
            .compile_expr(context, block_ref, &*var_decl_expr.init_val)?
            .unwrap();

        if !var_decl_expr.r#type.shape.is_empty() {
            let results = &[self.get_type(context, &var_decl_expr.r#type.shape)?];
            let operands = &[value];
            let reshape_op = OperationBuilder::new("toy.reshape", location)
                .add_operands(operands)
                .add_results(results)
                .build()?;
            let reshape_op_ref = block_ref.append_operation(reshape_op);
            value = reshape_op_ref.result(0).unwrap().into();
        }

        assert!(mlir_declare(
            &self.scoped_hash_table,
            &var_decl_expr.name,
            value
        ));
        Ok(())
    }

    fn compile_return_expr(
        &self,
        context: &Context,
        block_ref: BlockRef,
        return_expr: &crate::ast::ReturnExpr,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let input = match &return_expr.expr {
            Some(expr) => vec![self.compile_expr(context, block_ref, &*expr)?.unwrap()],
            None => Vec::new(),
        };
        let return_op = OperationBuilder::new("toy.return", location)
            .add_operands(&input)
            .build()?;
        block_ref.append_operation(return_op.into());
        Ok(())
    }

    fn compile_binary_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        binary_expr: &crate::ast::BinaryExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let lhs = Box::new(self.compile_expr(context, block_ref, &*binary_expr.lhs)?).unwrap();
        let rhs = Box::new(self.compile_expr(context, block_ref, &*binary_expr.rhs)?).unwrap();
        let op_ref = match binary_expr.op {
            crate::ast::BinaryOperator::Add => {
                let results = &[self.get_type(context, &[])?];
                let add_op = OperationBuilder::new("toy.add", location)
                    .add_results(results)
                    .add_operands(&[lhs, rhs])
                    .build()?;
                block_ref.append_operation(add_op.into())
            }
            crate::ast::BinaryOperator::Mul => {
                let results = &[self.get_type(context, &[])?];
                let mul_op = OperationBuilder::new("toy.mul", location)
                    .add_results(results)
                    .add_operands(&[lhs, rhs])
                    .build()?;
                block_ref.append_operation(mul_op.into())
            }
        };
        Ok(op_ref.result(0).unwrap().into())
    }

    fn compile_call_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        call_expr: &crate::ast::CallExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let callee = call_expr.callee.clone();
        let op_ref = if callee == "transpose" {
            assert!(call_expr.args.len() == 1);

            let input = self
                .compile_expr(context, block_ref, &call_expr.args[0])?
                .unwrap();
            let results = &[self.get_type(context, &[])?];
            let operands = &[input];
            let op = OperationBuilder::new("toy.transpose", location)
                .add_results(results)
                .add_operands(operands)
                .build()?;

            block_ref.append_operation(op)
        } else {
            let function = FlatSymbolRefAttribute::new(context, &callee);
            let mut arguments = Vec::new();
            for arg in &call_expr.args {
                arguments.push(self.compile_expr(context, block_ref, arg)?.unwrap());
            }
            let result_types = &[self.get_type(context, &[])?];
            let op = OperationBuilder::new("toy.generic_call", location)
                .add_attributes(&[(Identifier::new(context, "callee"), function.into())])
                .add_operands(&arguments)
                .add_results(result_types)
                .build()?;
            block_ref.append_operation(op.into())
        };
        let result = op_ref.result(0).unwrap().into();
        Ok(result)
    }

    fn compile_print_expr(
        &self,
        context: &Context,
        block_ref: BlockRef,
        print_expr: &crate::ast::PrintExpr,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let input = self
            .compile_expr(context, block_ref, &print_expr.arg)?
            .unwrap();
        let print_op = OperationBuilder::new("toy.print", location)
            .add_operands(&[input])
            .build()?;
        block_ref.append_operation(print_op.into());
        Ok(())
    }

    fn collect_expr_values(&self, expr: &crate::ast::Expr, values: &mut Vec<f64>) -> Result<()> {
        match expr {
            crate::ast::Expr::NumberExpr(number_expr) => values.push(number_expr.val),
            crate::ast::Expr::LiteralExpr(literal_expr) => {
                self.collect_literal_expr_values(literal_expr, values)?
            }
            _ => todo!(),
        }
        Ok(())
    }

    fn collect_literal_expr_values(
        &self,
        literal_expr: &crate::ast::LiteralExpr,
        values: &mut Vec<f64>,
    ) -> Result<()> {
        for expr in &literal_expr.values {
            self.collect_expr_values(expr, values)?
        }
        Ok(())
    }
}
